import dgl

# 假设已经加载了图
graphs = dgl.load_graphs('/home/hujiahui/git/paper/graph/timing_gcn/aes_cipher.graph.bin')
graph = graphs[0][0]  # 选择第一个图

# 获取编号为2的节点
node_id = 2

# 定义三种边类型
edge_types = ['net_in', 'net_out', 'cell_out']

# 遍历每个边类型，获取出边和入边的邻居
for edge_type in edge_types:
    print(f"\n查询边类型 {edge_type} 的邻居节点：")
    
    # 获取出边邻居节点
    out_neighbors = graph.successors(node_id, etype=edge_type)
    
    # 获取入边邻居节点
    in_neighbors = graph.predecessors(node_id, etype=edge_type)

    # 打印出与节点2相连的邻居节点
    print(f"节点 {node_id} 的出边邻居节点：")
    for neighbor in out_neighbors:
        print(f"邻居节点 {neighbor} 连接类型: {edge_type}")

    print(f"\n节点 {node_id} 的入边邻居节点：")
    for neighbor in in_neighbors:
        print(f"邻居节点 {neighbor} 连接类型: {edge_type}")
